package com.blyat.xsoft.kernel.component.security.oauth2.token.jwt;

import com.blyat.xsoft.kernel.component.security.oauth2.OauthConstant;
import com.blyat.xsoft.kernel.component.security.oauth2.token.OauthAccessTokenConverter;
import com.blyat.xsoft.kernel.util.SM4Util;
import com.google.common.base.Strings;
import org.springframework.security.core.userdetails.UserDetailsService;
import org.springframework.security.jwt.JwtHelper;
import org.springframework.security.oauth2.common.*;
import org.springframework.security.oauth2.common.util.JsonParser;
import org.springframework.security.oauth2.common.util.JsonParserFactory;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.store.JwtAccessTokenConverter;

import java.util.Date;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Objects;

/**
 * @Auther: syh
 * @Date: 2020/7/16
 * @Description:
 */
public class OauthJwtAccessTokenConverter extends JwtAccessTokenConverter {
    private JsonParser objectMapper = JsonParserFactory.create();

    public OauthJwtAccessTokenConverter(UserDetailsService userService) {
        // 使用SecurityContextHolder.getContext().getAuthentication()能获取到User信息
        super.setAccessTokenConverter(new OauthAccessTokenConverter(userService));
    }

    @Override
    public OAuth2AccessToken enhance(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
        DefaultOAuth2AccessToken result = new DefaultOAuth2AccessToken(accessToken);
        Map<String, Object> info = new LinkedHashMap<String, Object>(accessToken.getAdditionalInformation());
        String tokenId = result.getValue();
        if (!info.containsKey(TOKEN_ID)) {
            info.put(TOKEN_ID, tokenId);
        } else {
            tokenId = (String) info.get(TOKEN_ID);
        }

        // access_token 包含自动刷新过期token需要的数据(client_id/secret/refresh_token)
        Map<String, Object> details = (Map<String, Object>) authentication.getUserAuthentication().getDetails();
        String oauthClientId = "", oauthClientSecret = "";
        if (!Objects.isNull(details) && details.size() > 0) {
            try {
                if (details.get("client_id") != null) {
                    oauthClientId = SM4Util.encryptEcb(details.get("client_id").toString());
                } else {
                    oauthClientId = details.get(OauthConstant.OAUTH_CLIENT_ID).toString();
                }
                info.put(OauthConstant.OAUTH_CLIENT_ID, oauthClientId);

                if (details.get("client_secret") != null) {
                    oauthClientSecret = SM4Util.encryptEcb(details.get("client_secret").toString());
                } else {
                    oauthClientSecret = details.get(OauthConstant.OAUTH_CLIENT_SECRET).toString();
                }
                info.put(OauthConstant.OAUTH_CLIENT_SECRET, oauthClientSecret);
            } catch (Exception e) {
            }
        }

        OAuth2RefreshToken refreshToken = result.getRefreshToken();
        if (refreshToken != null) {
            DefaultOAuth2AccessToken encodedRefreshToken = new DefaultOAuth2AccessToken(accessToken);
            encodedRefreshToken.setValue(refreshToken.getValue());
            // Refresh tokens do not expire unless explicitly of the right type
            encodedRefreshToken.setExpiration(null);
            try {
                Map<String, Object> claims = objectMapper
                        .parseMap(JwtHelper.decode(refreshToken.getValue()).getClaims());
                if (claims.containsKey(TOKEN_ID)) {
                    encodedRefreshToken.setValue(claims.get(TOKEN_ID).toString());
                }
            } catch (IllegalArgumentException e) {
            }
            Map<String, Object> refreshTokenInfo = new LinkedHashMap<String, Object>(
                    accessToken.getAdditionalInformation());
            refreshTokenInfo.put(TOKEN_ID, encodedRefreshToken.getValue());
            // refresh token包含client id/secret, 自动刷新过期token时用到。
            if (!Strings.isNullOrEmpty(oauthClientId)) {
                refreshTokenInfo.put(OauthConstant.OAUTH_CLIENT_ID, oauthClientId);
            }
            if (!Strings.isNullOrEmpty(oauthClientSecret)) {
                refreshTokenInfo.put(OauthConstant.OAUTH_CLIENT_SECRET, oauthClientSecret);
            }
            refreshTokenInfo.put(ACCESS_TOKEN_ID, tokenId);
            encodedRefreshToken.setAdditionalInformation(refreshTokenInfo);
            DefaultOAuth2RefreshToken token = new DefaultOAuth2RefreshToken(
                    encode(encodedRefreshToken, authentication));
            if (refreshToken instanceof ExpiringOAuth2RefreshToken) {
                Date expiration = ((ExpiringOAuth2RefreshToken) refreshToken).getExpiration();
                encodedRefreshToken.setExpiration(expiration);
                token = new DefaultExpiringOAuth2RefreshToken(encode(encodedRefreshToken, authentication), expiration);
            }
            result.setRefreshToken(token);
            try {
                info.put(OauthConstant.OAUTH_REFRESH_TOKEN, SM4Util.encryptEcb(token.getValue()));
            } catch (Exception e) {}
        }
        result.setAdditionalInformation(info);
        result.setValue(encode(result, authentication));
        return result;
    }
}
